#include "registration.h"
#include <pcl/registration/icp_nl.h>
#include <pcl/registration/icp.h>
#include <pcl/registration/transformation_estimation_svd.h>
#include <pcl/registration/transformation_estimation_lm.h>
#include <pcl/registration/transformation_validation.h>
#include <pcl/registration/transformation_validation_euclidean.h>
#include <pcl/filters/voxel_grid.h>
#include <pcl/kdtree/impl/kdtree_flann.hpp>
#include <pcl/kdtree/kdtree_flann.h>
#include <pcl/keypoints/uniform_sampling.h>
#include <pcl/search/kdtree.h>
#include <pcl/features/fpfh_omp.h>
#include <pcl/features/normal_3d.h>
#include <pcl/features/normal_3d_omp.h>
#include <pcl/features/pfh.h>
#include <pcl/features/vfh.h>
#include <pcl/search/kdtree.h>
#include <pcl/filters/filter.h>
#include <math.h>
#include <utility>
#include <pcl/registration/correspondence_estimation.h>
#include <pcl/registration/correspondence_rejection_sample_consensus.h>
#include <pcl/features/normal_3d.h>



void removeNaNFromFeatures (PointCloud<FPFHSignature33>::Ptr features)
{
    PointCloud<FPFHSignature33>::Ptr filteredFeatures (new PointCloud<FPFHSignature33>);
    bool isNaN=false;
    for (int i=0;i<features->size();i++)
    {
        for (int j=0;j<33;j++)//scorre tutta la feature
            if (pcl_isnan (features->at(i).histogram[1]))
            {//se trova un NaN imposta il flag true..
                isNaN=true;
                break;
            }
        if (!isNaN)
            filteredFeatures->push_back(features->at(i));
    }
    *features=*filteredFeatures;
}

void onlyNormals(PointCloud<PointXYZ>::Ptr cloud,PointCloud<Normal>::Ptr normals, double radius)
{
    cout << "Normals estimation... " << flush;
    NormalEstimationOMP<PointXYZ, Normal> normal_estimation;
    normal_estimation.setSearchMethod (search::Search<PointXYZ>::Ptr (new search::KdTree<PointXYZ>));
    normal_estimation.setRadiusSearch (radius);
    normal_estimation.setInputCloud (cloud);//input, ho messo kpts per accelerare
    normal_estimation.compute (*normals);
    cout << "OK! Normals found: " << normals->points.size() << endl;
}

void fpfhWithUniformKeypoints(PointCloud<PointXYZ>::Ptr source, PointCloud<PointXYZI>::Ptr keypoints, PointCloud<FPFHSignature33>::Ptr features)
{
    PointCloud<PointXYZ>::Ptr kpts(new PointCloud<PointXYZ>);
    copyPointCloud(*keypoints, *kpts);//passo da XYZI a XYZ
    PointCloud<Normal>::Ptr normals (new  PointCloud<Normal>);
    onlyNormals(kpts,normals,20);//20
    cout << "FPFH Descriptors extraction... " << flush;
    FPFHEstimationOMP<PointXYZ, Normal, FPFHSignature33>::Ptr feature_extractor (new FPFHEstimationOMP<PointXYZ, Normal, FPFHSignature33>);
    feature_extractor->setSearchMethod (search::Search<PointXYZ>::Ptr (new search::KdTree<PointXYZ>));
    feature_extractor->setRadiusSearch (25);//25
    feature_extractor->setInputNormals(normals);
    feature_extractor->setInputCloud(kpts);
    feature_extractor->compute (*features);
    removeNaNFromFeatures(features);
    //printFeatures(features);

    cout << "OK! Features found: " << features->points.size() << endl;

}

float computeLeafSize (PointCloud<PointXYZ>::Ptr cloud,int nVoxelsRequired)
{
    PointXYZ min;
    PointXYZ max;
    //trova gli estremi della cloud
    getMinMax3D<PointXYZ>(*cloud,min,max);
    Eigen::Vector3f size=max.getVector3fMap()-min.getVector3fMap();
    float surface=size[0]*size[1];
    float leafSize= sqrt(2*surface/nVoxelsRequired);
    return leafSize;
}

void uniformKeypoints(PointCloud<PointXYZ>::Ptr input,PointCloud<PointXYZI>::Ptr keypoints, double leafSize)
{
    cout << "Uniform Keypoint detection..." << flush;
    UniformSampling<PointXYZ>* us= new UniformSampling<PointXYZ>;
    search::KdTree<PointXYZ>::Ptr tree (new search::KdTree<PointXYZ> );
    us->setInputCloud(input);
    us->setRadiusSearch(leafSize);
    us->setSearchMethod(tree);
    PointCloud<int>::Ptr temp(new PointCloud<int>);
    us->compute(*temp);
    for (int i=0;i<temp->size();i++)
    {
        PointXYZI point;
        point.x=input->at( temp->at(i)).x;
        point.y=input->at( temp->at(i)).y;
        point.z=input->at( temp->at(i)).z;
        keypoints->push_back(point);
    }
    cout << "OK! Keypoints found: " << keypoints->points.size() << endl;
}

void voxelKeypoints (PointCloud<PointXYZ>::Ptr input, PointCloud<PointXYZI>::Ptr keypoints, double leafSize)
{
    cout << "Voxelling and keypoints... "<<flush; //DEBUG
    PointCloud<PointXYZ>::Ptr cloud_filtered (new PointCloud<PointXYZ>);
    VoxelGrid<PointXYZ> sor;
    sor.setInputCloud (input);
    sor.setLeafSize (leafSize,leafSize,leafSize);
    sor.filter (*cloud_filtered);
    copyPointCloud(*cloud_filtered,*keypoints);
    cout << "Voxel keypoints: " << keypoints->points.size()<< endl;
}


void findCorrespondencesFPFH (
        PointCloud<FPFHSignature33>::Ptr source,
        PointCloud<FPFHSignature33>::Ptr target,
        CorrespondencesPtr correspondences)
{
    cout << "Finding correspondences..." << flush;

    registration::CorrespondenceEstimation<FPFHSignature33,FPFHSignature33> ce;
    ce.setInputCloud(source);
    ce.setInputTarget(target);
    ce.determineReciprocalCorrespondences(*correspondences);//evidentemente non va senza reciprocal
    cout << "OK! Correspondences found: " << correspondences->size() << endl;
}


void filterCorrespondencesSAC ( PointCloud<PointXYZI>::Ptr source_keypoints,
                                PointCloud<PointXYZI>::Ptr target_keypoints,
                                CorrespondencesPtr correspondences,
                                double inlierThreshold)
{
    //cout << "Filtering correspondences SAC..." << flush;
    registration::CorrespondenceRejectorSampleConsensus<PointXYZI> rejector;
    rejector.setInputCloud(source_keypoints);
    rejector.setTargetCloud(target_keypoints);
    rejector.setMaxIterations(50);
    rejector.setInlierThreshold(inlierThreshold);
    rejector.setInputCorrespondences(correspondences);
    rejector.getCorrespondences(*correspondences);
    //cout << "OK! Correspondences filtered: " << correspondences->size()<< endl;
}

void initialRegisterSourceToTarget (PointCloud<PointXYZ>::Ptr source, PointCloud<PointXYZ>::Ptr target, PointCloud<PointXYZ>::Ptr registered,
                                    PointCloud<PointXYZI>::Ptr source_keypoints, PointCloud<PointXYZI>::Ptr target_keypoints,
                                    CorrespondencesPtr correspond, Eigen::Matrix4f &initial_transformation_matrix, float resolution)
{
    time_t start,end;
    time (&start);

    cout << "Initial Leaf Size: " << resolution<<endl;
    voxelKeypoints (source, source_keypoints, resolution); voxelKeypoints(target, target_keypoints, resolution);



    //FPFH
    PointCloud<FPFHSignature33>::Ptr source_features (new PointCloud<FPFHSignature33>);
    PointCloud<FPFHSignature33>::Ptr target_features (new PointCloud<FPFHSignature33>);
    fpfhWithUniformKeypoints (source,source_keypoints, source_features);
    fpfhWithUniformKeypoints (target,target_keypoints, target_features);


    findCorrespondencesFPFH(source_features,target_features,correspond);

    double thresh=5;
    float bestScore=FLT_MAX;
    double bestThresh=0;
    CorrespondencesPtr currentCorrespond(new Correspondences);
    CorrespondencesPtr bestCorrespond(new Correspondences);

    //TRANSFORM ESTIMATION
    registration::TransformationEstimationSVD<PointXYZI, PointXYZI>::Ptr transformation_estimation (new registration::TransformationEstimationSVD<PointXYZI, PointXYZI>);

    while(thresh<20)
    {
        *currentCorrespond=*correspond;
        //FILTER CORRESPONDENCES
        filterCorrespondencesSAC(source_keypoints,target_keypoints,currentCorrespond,thresh);
        //cout << "Initial matrix calculation with thresh "<< thresh <<"..." << flush;
        transformation_estimation->estimateRigidTransformation (*source_keypoints, *target_keypoints, *currentCorrespond, initial_transformation_matrix);
        PointCloud<PointXYZI>::Ptr registered_keypoints(new PointCloud<PointXYZI>);
        transformPointCloud(*source_keypoints, *registered_keypoints, initial_transformation_matrix);
        float score=validationScoreProximityTarget<PointXYZI>(*registered_keypoints,*target_keypoints);
        //cout << " score: "<< score <<flush;
        if(score<bestScore)
        {
            bestScore=score;
            bestThresh=thresh;
            *bestCorrespond= *currentCorrespond;
        }
        thresh+=1;
    }
    *correspond=*bestCorrespond;
    cout << "Initial matrix calculation..." << flush;
    transformation_estimation->estimateRigidTransformation (*source_keypoints, *target_keypoints, *correspond, initial_transformation_matrix);
    PointCloud<PointXYZI>::Ptr registered_keypoints(new PointCloud<PointXYZI>);
    transformPointCloud(*source_keypoints, *registered_keypoints, initial_transformation_matrix);
    transformPointCloud(*source, *registered, initial_transformation_matrix);
    cout << "OK! Used "<< correspond->size() <<" corresponcences. Score on keypoints: " << bestScore<<endl;


    time (&end);
    double dif = difftime (end,start);
    cout << "------------------------------------->  Initial registration in " << dif << " seconds"<< endl;
}


void
finalRegisterSourceToTarget (
        PointCloud<PointXYZ>::Ptr source ,
        PointCloud<PointXYZ>::Ptr target,
        PointCloud<PointXYZ>::Ptr registered,
        Eigen::Matrix4f &final_transformation_matrix,
        float resolution)
{
    time_t start,end;
    time (&start);
 cout << "Final Leaf Size: " << resolution<<endl;
    PointCloud<PointXYZI>::Ptr source_keypoints (new PointCloud<PointXYZI>);
    PointCloud<PointXYZI>::Ptr source_keypoints_registered (new PointCloud<PointXYZI>);
    PointCloud<PointXYZI>::Ptr target_keypoints (new PointCloud<PointXYZI>);

    //uniformKeypoints(source,source_keypoints,resolution); uniformKeypoints(target,target_keypoints,resolution);
    voxelKeypoints(source,source_keypoints,resolution); voxelKeypoints(target,target_keypoints,resolution);

    //agisco sui keypoints, poi applico alla cloud originale
    cout << "SourceKpts: "<<source_keypoints->size()<<" TargetKpts: "<<target_keypoints->size()<<" Final matrix calculation..." << flush;
    //----------------------
    IterativeClosestPoint<PointXYZI, PointXYZI>::Ptr registration (new IterativeClosestPoint<PointXYZI, PointXYZI>);
    //IterativeClosestPointNonLinear<PointXYZI, PointXYZI>::Ptr registration (new IterativeClosestPointNonLinear<PointXYZI, PointXYZI>);
    //----------------------
    registration->setInputCloud(source_keypoints);
    registration->setInputTarget (target_keypoints);
    //registration->setMaxCorrespondenceDistance(50); //4
    registration->setRANSACOutlierRejectionThreshold (10);//0.05 default  Set the inlier distance threshold for the internal RANSAC outlier rejection loop.
    //The method considers a point to be an inlier, if the distance between the target data index and the transformed source index is smaller than the given inlier distance threshold. The value is set by default to 0.05m.
    registration->setRANSACIterations(100);//default 1000
    //----------------------------------------
    //registration->setEuclideanFitnessEpsilon(1);  //1 TermCriterion  The sum of Euclidean squared errors is smaller than a user defined threshold
    registration->setMaximumIterations (100);     //2 TermCriterion  Number of iterations has reached the maximum
    registration->setTransformationEpsilon (1e-7);//3 TermCriterion  The epsilon (difference) between the previous transformation and the current estimated transformation is smaller than parameter
    //----------------------------------------
    registration->align(*source_keypoints_registered);
    final_transformation_matrix = registration->getFinalTransformation();
    if (registration->hasConverged()) cout << " Converged! ";
    cout << "OK! Fitness Score: " << validationScoreProximityTarget<PointXYZI>(*source_keypoints_registered,*target_keypoints)  << endl;
    transformPointCloud(*source, *registered, final_transformation_matrix);

    time (&end);
    double dif = difftime (end,start);
    cout << "------------------------------------->  Final registration in " << dif << " seconds"<< endl;
}


void
registerSourceToTarget (
        PointCloud<PointXYZRGB>::Ptr sourceColor ,
        PointCloud<PointXYZRGB>::Ptr targetColor,
        PointCloud<PointXYZRGB>::Ptr registeredColor)
{
    time_t start,end;
    time (&start);


    //CANCELLA WARNING PCL
    //console::setVerbosityLevel(console::L_ALWAYS);

    PointCloud<PointXYZ>::Ptr source(new PointCloud<PointXYZ>);
    PointCloud<PointXYZ>::Ptr target(new PointCloud<PointXYZ>);

    copyPointCloud(*sourceColor,*source);
    copyPointCloud(*targetColor,*target);


    PointCloud<PointXYZ>::Ptr initialRegistered(new PointCloud<PointXYZ>);
    PointCloud<PointXYZ>::Ptr registered(new PointCloud<PointXYZ>);
    PointCloud<PointXYZI>::Ptr source_keypoints (new PointCloud<PointXYZI> );
    PointCloud<PointXYZI>::Ptr target_keypoints (new PointCloud<PointXYZI> );
    CorrespondencesPtr correspond(new Correspondences);



    //REGISTRATION
    float leafSize=computeLeafSize(target,4000);
    float initialResolution=leafSize*1.0;
    float finalResolution=leafSize*0.5;


    Eigen::Matrix4f initial_transformation_matrix = Eigen::Matrix4f::Identity();
    Eigen::Matrix4f final_transformation_matrix = Eigen::Matrix4f::Identity();
    initialRegisterSourceToTarget(source,target,initialRegistered,source_keypoints,target_keypoints,correspond,initial_transformation_matrix,initialResolution);
    transformPointCloud(*sourceColor,*registeredColor,initial_transformation_matrix);


    finalRegisterSourceToTarget(initialRegistered,target,registered,final_transformation_matrix,finalResolution);
    transformPointCloud(*registeredColor,*registeredColor,final_transformation_matrix);


    time (&end);
    double dif = difftime (end,start);
    cout << "------------------------------------->  Total registration in " << dif << " seconds"<< endl;
}



